import numpy as np
import matplotlib.pyplot as plt

# -----------------------------------------------------------------------------
# Práctica FTCS
tiempo = 10; particulas = 20 # Número de intervalos y demás datos, las particulas seran N+1
delta_t = 0.01; delta_x = 0.5; alpha = 1
iteracciones = int(tiempo/delta_t)  # Iteraciones que habrá (tirando a lo bajo)

T = np.zeros(particulas+1) #Número de puntos

T[9:11] = 10    # Condiciones Iniciales

Tnew = np.zeros(particulas+1)

plt.close('all')        
plt.figure(1); plt.plot(T,"-r"); plt.pause(0.1)

for n in range(iteracciones):   # Tic tac del reloj interno
    for i in range(1,particulas):   # Para cada particula
        Tnew[i]=T[i]+alpha*delta_t/delta_x/delta_x*(T[i+1]-2*T[i]+T[i-1])
        
    # Condiciones de frontera
    Tnew[0] = 5
    Tnew[particulas] = 3
    
    T[:] = Tnew[:]          # Actualiza Variables, ahora T es lo que era Tnew
    if n%10==0: # Graficación por cada 10 iteraciones
        plt.figure(1)
        plt.plot(T)
        #plt.ylim(0, 7) # Gráfico constante
        plt.pause(0.01)
        plt.show()
        
plt.xlabel("i-indice")
plt.ylabel("Temp")

# Ejercicio 1 -----------------------------------------------------------------

def FTCS(N, tiempo, delta_t, delta_x, alpha, init_func, cf_func):
    
    # Calculamos r (para ser estable debe ser menor a 1/2)
    r = alpha * delta_t / delta_x**2
    print(f"r ={r:.2f}")
    if r>1/2:
        print("r inestable. Aumenta delta_x o reduce delta_t")
    else:
        pass
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)    # Y formamos un array identico a T pero con todo ceros
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = T[i] + r * (T[i+1] - 2*T[i] + T[i-1])
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion, que puede necesitar n o dt (como sen())
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    """
            plt.plot(T)
    plt.xlabel("i")
    plt.ylabel("T")
    plt.show()
    """
    
    return T

# Funcion donde todos los valores iniciales son 0
def init_None(N):
    T = np.zeros(N+1)
    return T

# Funcion donde todos los valores iniciales son random entre 0 y 10
def init_random(N):
    T = np.random.rand(N+1)*10  # Valores entre 0 y 10
    return T

# Las fronteras son fijas, focos fijos en los lados (0 en la izquierda, 10 en la derecha)
def dirichlet(T, n, dt):
    T[0] = 0
    T[-1] = 10
    return T

# Funcion rara donde la frontera izquierda es siempre 0 y la derecha va oscilando
def seno(T, n, dt):
    t = n * dt
    T[0] = 0
    T[-1] = np.sin(0.5 * t)
    return T

# Condiciones de frontera de flujo nulo, como si fuera una piscina, los bordes van moviendose
def flujo_nulo(T, n, dt):
    T[0] = T[1]
    T[-1] = T[-2]
    return T

# Apartado a)
FTCS(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_None, cf_func=dirichlet)

# Apartado b)
FTCS(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_random, cf_func=dirichlet)

# Apartado c)
FTCS(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_None, cf_func=seno)

# Apartado d)
FTCS(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_None, cf_func=dirichlet)

# Ejercicio 2 -----------------------------------------------------------------

def esquema_tres_niveles_temporales(N, tiempo, delta_t, delta_x, alpha, init_func, cf_func):
    
    # Calculamos r (ahora no cumple la misma condicion de estabilidad que FTCS de ser menor a 1/2)
    r = alpha * delta_t / delta_x**2

    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)
    Told = np.copy(T)
    # Para salvar situacion, avanzar un tiempo aplicando FTCS que se aproxima a nuestro modelo como arranque
    T_primer_paso = np.copy(T)  # Primer paso de FTCS aproximado y luego ya lo normal y Tn dirige
    for i in range(1, N):
        T_primer_paso[i] = T[i] + r * (T[i+1] - 2*T[i] + T[i-1])
    
    T_primer_paso = cf_func(T_primer_paso, 0, delta_t)   # COndiciones de frontera para el momento init, n=0
    T = np.copy(T_primer_paso) # Ahora T pasa a ser tiempo 1 y Told es tiempo 0
    # QUizas, probablemente se puede hacer sin el copy, pero asi queda mas seguro
    
    # Ploteamos este primer momento
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):
        
        for i in range(1, N):
            Tnew[i] = (2/3)*(-0.5*Told[i] + 2*T[i] + r*(T[i+1] + T[i-1] - 2*T[i]))  # Funcion especifica que calcula el Tnew
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion que puede necesitar n o dt (como sen())
        
        Told[:] = T[:]
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion, primero T old desde T y luego T desde T new
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos cada 10 iteraciones
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    return T

# Apartado a)
esquema_tres_niveles_temporales(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_None, cf_func=dirichlet)

# Apartado b)
esquema_tres_niveles_temporales(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_random, cf_func=flujo_nulo)

# Ejercicio 3 -----------------------------------------------------------------

def cinco_puntos(N, tiempo, delta_t, delta_x, alpha, init_func, cf_func):
    
    r = alpha * delta_t / delta_x**2    # Da igual el r aqui tambien
    print(f"r ={r:.2f}")
    
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)
    
    # Ploteamos el primer momento
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):
        
        for i in range(1, N):
            
            if i == 1 or i == N-1: # FTCS en los extremos
                Tnew[i] = T[i] + r * (T[i+1] - 2*T[i] + T[i-1])
            else:
                Tnew[i] = r* ((-1/12)*T[i-2] + (4/3)*T[i-1] -2.5*T[i] + (4/3)*T[i+1] + (-1/12)*T[i+2]) +T[i]
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion que puede necesitar n o dt (como sen())
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    return T

# Apartado a)
cinco_puntos(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_None, cf_func=dirichlet)

# Apartado b)
cinco_puntos(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_random, cf_func=dirichlet)

# Ejercicio 4 -----------------------------------------------------------------

def dufort_frankel(N, tiempo, delta_t, delta_x, alpha, init_func, cf_func):
    
    # Calculamos r (ahora no cumple la misma condicion de estabilidad que FTCS de ser menor a 1/2)
    r = alpha * delta_t / delta_x**2

    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Told = np.copy(T)
    Tnew = np.zeros(N+1)
    # Para salvar situacion, avanzar un tiempo aplicando FTCS que se aproxima a nuestro modelo como arranque
    T_primer_paso = np.copy(T)  # Primer paso de FTCS aproximado y luego ya lo normal y Tn dirige
    for i in range(1, N):
        T_primer_paso[i] = T[i] + r * (T[i+1] - 2*T[i] + T[i-1])
    
    T_primer_paso = cf_func(T_primer_paso, 0, delta_t)   # COndiciones de frontera para el momento init, n=0
    T = np.copy(T_primer_paso) # Ahora T pasa a ser tiempo 1 y Told es tiempo 0
    # QUizas se puede hacer sin el copy, pero asi queda mas seguro
    
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):
        
        for i in range(1, N):
            Tnew[i] = (1/(1+2*r))*(Told[i] + 2*r*(T[i-1]+T[i+1]-Told[i]))
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion que puede necesitar n o dt (como sen())
        
        Told[:] = T[:]
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    return T

# Apartado a)
dufort_frankel(N=20, tiempo=5, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_None, cf_func=dirichlet)

# Apartado b)
dufort_frankel(N=20, tiempo=5, delta_t=0.01, delta_x=0.1, alpha=1/4, init_func=init_random, cf_func=dirichlet)

# Ejercicio 8 -----------------------------------------------------------------
def condiciones_centro(T, N):
    T[N//2] = 5
    T[N//2 - 1] = 5
    
    return T
    
def condiciones_izquierda(Tnew, n, dt, X):
    x, y = X
    b = 5
    
    dxdt = y
    dydt = -b*y-x
    
    x = x + dt*dxdt
    y = y + dt*dydt
    
    X = np.array([x, y])
    
    Tnew[0] = x
    
    return Tnew, X
    
def condiciones_derecha(T):
    T[-1] = T[-2]
    
    return T

def FTCS_ultimoej(N, tiempo, delta_t, delta_x, alpha, init_func):
    
    r = alpha * delta_t / delta_x**2
    print(f"r ={r:.2f}")
    if r>1/2:
        print("r inestable. Aumenta delta_x o reduce delta_t")
    else:
        pass
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)    # Y formamos un array identico a T pero con todo ceros
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    X = np.array([1.5, 0.1])    # Condiciones iniciales de la ecuacion diferencial
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = T[i] + r * (T[i+1] - 2*T[i] + T[i-1])
        
        # Condiciones de frontera externas y centrales
        Tnew = condiciones_centro(Tnew, N)
        Tnew, X = condiciones_izquierda(Tnew, n, delta_t, X)
        Tnew = condiciones_derecha(Tnew)
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i")
            plt.ylabel("T")
    plt.show()
    
    return T

FTCS_ultimoej(20, 25, 0.1, 0.5, 1/4, init_None)